#include "Socket.h"

int Socket::nofSockets_= 0;

void Socket::Start() 
{
	if (!nofSockets_) 
	{
		WSADATA info;
		if (WSAStartup(MAKEWORD(2,0), &info))
		{
			throw "Could not start WSA";
		}
	}
	++nofSockets_;
}

void Socket::End() 
{
	WSACleanup();
}

Socket::Socket( OPTIONAL SOCKET _socket /* = NULL */, OPTIONAL UINT _socketType /* = SOCK_STREAM */, OPTIONAL sockaddr *_socketAddr /* = NULL */ ) 
: s_(_socket)
{
	Start();

	if (NULL == s_)
	{
		if ( SOCK_DGRAM == _socketType)
		{
			s_ = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
		}
		else
		{
			s_ = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		}
	}

	if (INVALID_SOCKET == s_)
	{
		throw "Invalid socket";
	}

	if (NULL != _socketAddr)
	{
		memcpy(s_addr_,_socketAddr,sizeof(sockaddr));
	}

	refCounter_ = new int(1);
}

Socket::~Socket() {
	if (! --(*refCounter_))
	{
		Close();
		delete refCounter_;
	}

	--nofSockets_;
	if (!nofSockets_) End();
}

Socket::Socket(const Socket& o)
{
	refCounter_=o.refCounter_;
	(*refCounter_)++;
	s_         =o.s_;

	nofSockets_++;
}

Socket& Socket::operator=(Socket& o)
{
	(*o.refCounter_)++;

	refCounter_=o.refCounter_;
	s_         =o.s_;

	nofSockets_++;

	return *this;
}

void Socket::Close()
{
	closesocket(s_);
}

void Socket::ReadBytes(IN LPBYTE bytes,IN OUT ULONG &count) 
{
	if ((NULL == bytes) || (count = 0))
	{
		count = 0;
		return;
	}

	ULONG myCount = count;
	ULONG totalRcv = 0;

	while (1) {
		int rv = recv (s_, (char*)bytes, myCount, 0);
		if (SOCKET_ERROR == rv)
		{
			// we have an error - do not send any data across
			memset(bytes,NULL,count);
			totalRcv = 0;
			break;
		}
		if (rv == 0) break;

		totalRcv += rv;
		myCount -= rv;

		if (myCount <= 0) break;
	}

	count = totalRcv;
}

void Socket::SendBytes(IN const LPBYTE bytes,IN ULONG count)
{
	if ((NULL == bytes) || (count <= 0))
		return;

	if (SOCKET_ERROR == send(s_,(char*)bytes,count,0))
	{
		throw "Error";
	}
}
